//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package callcmd

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sync/atomic"
	"time"

	"atomgit.com/openlinksaas-org/grpcutil/utils"
	"github.com/golang/protobuf/jsonpb" //lint:ignore SA1019 we have to import this because it appears in exported API
	"github.com/golang/protobuf/proto"  //lint:ignore SA1019 we have to import this because it appears in exported API
	"github.com/jhump/protoreflect/desc"
	"github.com/jhump/protoreflect/dynamic"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/types/descriptorpb"
)

var inputFile = ""
var remoteAddr = ""
var secure = false
var serviceName = ""
var methodName = ""

var readTime = int64(0)

var CallCommand = &cobra.Command{
	Use: "call",
	RunE: func(cmd *cobra.Command, args []string) error {
		readTime = time.Now().Unix()
		return runCallCommand()
	},
}

func runCallCommand() error {
	data, err := os.ReadFile(inputFile)
	if err != nil {
		return err
	}
	descSet := &descriptorpb.FileDescriptorSet{}
	err = proto.Unmarshal(data, descSet)
	if err != nil {
		return err
	}

	unresolved := map[string]*descriptorpb.FileDescriptorProto{}
	for _, fd := range descSet.File {
		unresolved[fd.GetName()] = fd
	}
	resolved := map[string]*desc.FileDescriptor{}
	for _, fd := range descSet.File {
		_, err := utils.ResolveFileDescriptor(unresolved, resolved, fd.GetName())
		if err != nil {
			return err
		}
	}
	registry := dynamic.NewExtensionRegistryWithDefaults()
	for _, fd := range resolved {
		registry.AddExtensionsFromFile(fd)
	}
	msgFactory := dynamic.NewMessageFactoryWithExtensionRegistry(registry)

	var reqType *desc.MessageDescriptor
	var respType *desc.MessageDescriptor
	clientStream := false
	serverStream := false
	for _, fd := range resolved {
		svc := fd.FindService(serviceName)
		if svc == nil {
			continue
		}
		method := svc.FindMethodByName(methodName)
		if method == nil {
			continue
		}
		reqType = method.GetInputType()
		respType = method.GetOutputType()
		clientStream = method.IsClientStreaming()
		serverStream = method.IsServerStreaming()
		break
	}
	if reqType == nil {
		return fmt.Errorf("unkown service or method")
	}
	req := msgFactory.NewDynamicMessage(reqType)
	resp := msgFactory.NewDynamicMessage(respType)

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	options := []grpc.DialOption{}
	if secure {
		creds := credentials.NewTLS(nil)
		options = append(options, grpc.WithTransportCredentials(creds))
	} else {
		options = append(options, grpc.WithTransportCredentials(insecure.NewCredentials()))
	}
	conn, err := grpc.DialContext(ctx, remoteAddr, options...)
	if err != nil {
		return err
	}
	defer conn.Close()

	//超时关闭stdin
	go func() {
		for {
			time.Sleep(1 * time.Second)
			lastTime := atomic.LoadInt64(&readTime)
			nowTime := time.Now().Unix()
			if (nowTime - lastTime) > 2 {
				os.Stdin.Close()
				break
			}
		}
	}()

	reqMethod := fmt.Sprintf("/%s/%s", serviceName, methodName)
	if clientStream && serverStream {
		return invokeBidi(ctx, conn, reqMethod, req, resp)
	} else if clientStream {
		return invokeClientStream(ctx, conn, reqMethod, req, resp)
	} else if serverStream {
		return invokeServerStream(ctx, conn, reqMethod, req, resp)
	} else {
		return invokeUnary(ctx, conn, reqMethod, req, resp)
	}
}

func invokeBidi(ctx context.Context, conn *grpc.ClientConn, reqMethod string, req *dynamic.Message, resp *dynamic.Message) error {
	sd := &grpc.StreamDesc{
		StreamName:    methodName,
		ServerStreams: true,
		ClientStreams: true,
	}
	cs, err := conn.NewStream(ctx, sd, reqMethod)
	if err != nil {
		return err
	}
	var sendErr atomic.Value

	go func() {
		dec := json.NewDecoder(os.Stdin)
		for {
			var msg json.RawMessage
			err := dec.Decode(&msg)
			if err == io.EOF {
				break
			}
			if err != nil {
				sendErr.Store(err)
				break
			}
			atomic.StoreInt64(&readTime, time.Now().Unix())
			req.Reset()
			err = jsonpb.UnmarshalString(string(msg), req)
			if err != nil {
				sendErr.Store(err)
				break
			}
			err = cs.SendMsg(req)
			if err == io.EOF {
				break
			}
			if err != nil {
				sendErr.Store(err)
				break
			}
		}
	}()

	marshaler := &jsonpb.Marshaler{
		EmitDefaults: true,
		Indent:       "  ",
	}
	for {
		resp.Reset()
		err = cs.RecvMsg(resp)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		jsonStr, err := marshaler.MarshalToString(resp)
		if err != nil {
			return err
		}
		fmt.Print(jsonStr)
	}

	err, ok := sendErr.Load().(error)
	if ok {
		return err
	}
	return nil
}

func invokeClientStream(ctx context.Context, conn *grpc.ClientConn, reqMethod string, req *dynamic.Message, resp *dynamic.Message) error {
	dec := json.NewDecoder(os.Stdin)
	sd := &grpc.StreamDesc{
		StreamName:    methodName,
		ServerStreams: false,
		ClientStreams: true,
	}
	cs, err := conn.NewStream(ctx, sd, reqMethod)
	if err != nil {
		return err
	}
	marshaler := &jsonpb.Marshaler{
		EmitDefaults: true,
		Indent:       "  ",
	}
	for {
		var msg json.RawMessage
		err := dec.Decode(&msg)
		if err == io.EOF {
			err2 := cs.CloseSend()
			if err2 != nil {
				return err2
			}
			err2 = cs.RecvMsg(resp)
			if err2 != nil {
				return err2
			}
			break
		}
		if err != nil {
			return err
		}
		atomic.StoreInt64(&readTime, time.Now().Unix())
		req.Reset()
		err = jsonpb.UnmarshalString(string(msg), req)
		if err != nil {
			return err
		}
		err = cs.SendMsg(req)
		if err == io.EOF {
			err2 := cs.CloseSend()
			if err2 != nil {
				return err2
			}
			err2 = cs.RecvMsg(resp)
			if err2 != nil {
				return err2
			}
			break
		}
		if err != nil {
			return err
		}
	}
	jsonStr, err := marshaler.MarshalToString(resp)
	if err != nil {
		return err
	}
	fmt.Print(jsonStr)
	return nil
}

func invokeServerStream(ctx context.Context, conn *grpc.ClientConn, reqMethod string, req *dynamic.Message, resp *dynamic.Message) error {
	dec := json.NewDecoder(os.Stdin)
	var msg json.RawMessage
	err := dec.Decode(&msg)
	if err != nil {
		return err
	}
	err = jsonpb.UnmarshalString(string(msg), req)
	if err != nil {
		return err
	}
	sd := &grpc.StreamDesc{
		StreamName:    methodName,
		ServerStreams: true,
		ClientStreams: false,
	}
	cs, err := conn.NewStream(ctx, sd, reqMethod)
	if err != nil {
		return err
	}
	err = cs.SendMsg(req)
	if err != nil {
		return err
	}
	err = cs.CloseSend()
	if err != nil {
		return err
	}
	marshaler := &jsonpb.Marshaler{
		EmitDefaults: true,
		Indent:       "  ",
	}
	for {
		resp.Reset()
		err = cs.RecvMsg(resp)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		jsonStr, err := marshaler.MarshalToString(resp)
		if err != nil {
			return err
		}
		fmt.Print(jsonStr)
	}
	return nil
}

func invokeUnary(ctx context.Context, conn *grpc.ClientConn, reqMethod string, req *dynamic.Message, resp *dynamic.Message) error {
	dec := json.NewDecoder(os.Stdin)
	var msg json.RawMessage
	err := dec.Decode(&msg)
	if err != nil {
		return err
	}
	err = jsonpb.UnmarshalString(string(msg), req)
	if err != nil {
		return err
	}
	err = conn.Invoke(ctx, reqMethod, req, resp)
	if err != nil {
		return err
	}
	marshaler := &jsonpb.Marshaler{
		EmitDefaults: true,
		Indent:       "  ",
	}
	jsonStr, err := marshaler.MarshalToString(resp)
	if err != nil {
		return err
	}
	fmt.Print(jsonStr)
	return nil
}

func init() {
	CallCommand.Flags().StringVar(&inputFile, "in", "", "proto parse info file")
	CallCommand.Flags().StringVar(&remoteAddr, "addr", "", "grpc server addr")
	CallCommand.Flags().BoolVar(&secure, "secure", false, "use grpc secure mode")
	CallCommand.Flags().StringVar(&serviceName, "service", "", "service name")
	CallCommand.Flags().StringVar(&methodName, "method", "", "method name")
}
